import cv2 as cv
import numpy as np
import os
import sys
from python_ai.common.xcommon import *

np.random.seed(1)

VER = 'v5.0'
INTERVAL = 1000 // 25
BASE_DIR, FILE_NAME = os.path.split(__file__)
SAVE_DIR = os.path.join(BASE_DIR, '_save', FILE_NAME, VER)
HISTORY = 500
ALPHA = 0.001

# VIDEO_SOURCE = r"C:\DL_DATA\dl2_DataSets\traffic_counter\traffic_1min.mp4"
path = '../../../../large_data/video/traffic/traffic_20s.avi'
VIDEO_SOURCE = os.path.join(BASE_DIR, path)
# IMAGE_DIR = r"C:\DL_DATA\dl2_DataSets\traffic_counter\out1"
IMAGE_DIR = os.path.join(SAVE_DIR, 'out1')
# out_path = './out'
out_path = os.path.join(SAVE_DIR, 'out')
filtered_video_path = os.path.join(out_path, 'filtered.avi')

bg_subtractor = cv.createBackgroundSubtractorMOG2(history=HISTORY, detectShadows=False)

print(f'Learning for {HISTORY} history pictures ...')
video = cv.VideoCapture(VIDEO_SOURCE)
cnt = 0
while True:
    ret, img = video.read()
    if not ret:
        print('Video is over')
        break
    else:
        cnt += 1
        if cnt > HISTORY:
            break
        bg_subtractor.apply(img, None, ALPHA)
if cnt < HISTORY:
    print(f'Warning: Frames not enough! need:{HISTORY}, actual:{cnt}')
video.release()
print('Learnt.')

video = cv.VideoCapture(VIDEO_SOURCE)
os.makedirs(os.path.split(filtered_video_path)[0], exist_ok=True)
fourcc = cv.VideoWriter_fourcc(*list('XVID'))
writer = cv.VideoWriter(filtered_video_path, fourcc, INTERVAL, (img.shape[1], img.shape[0]))
cnt = 0
while True:
    ret, img = video.read()
    if not ret:
        print('Video is over')
        break
    else:
        img_ = img.copy()

        mask = bg_subtractor.apply(img, None, ALPHA)
        mask_ = mask.copy()

        # process mask
        # kernel = cv.getStructuringElement(cv.MORPH_ELLIPSE, (2, 2))
        kernel = np.ones((3, 3), dtype=np.uint8)
        mask = cv.morphologyEx(mask, cv.MORPH_OPEN, kernel)
        mask = cv.morphologyEx(mask, cv.MORPH_CLOSE, kernel)
        mask = cv.dilate(mask, kernel, iterations=2)

        img = cv.bitwise_and(img, img, mask=mask)
        writer.write(img)

        mask4draw = np.zeros_like(img, dtype=np.uint8)
        mask4draw[mask_ > 127] = 255.
        maskOri4Draw = np.zeros_like(img, dtype=np.uint8)
        maskOri4Draw[mask > 127] = 255.

        row01 = np.concatenate([img_, mask4draw], axis=1)
        row02 = np.concatenate([maskOri4Draw, img], axis=1)
        final = np.concatenate([row01, row02], axis=0)

        final = imzoom_rect(final, (1600*0.8, 900*0.8))
        cv.imshow('Traffic', final)
        cnt += 1
    k = cv.waitKey(INTERVAL) & 0xFF
    if 27 == k:
        print('Interrupted!')
        break
print(f'Frames count: {cnt}')
print(f'Written filtered video to {filtered_video_path}')
video.release()
writer.release()
cv.destroyAllWindows()
print('Over')
